import configs
from scripts.server.kb_api import get_chunks
from scripts.server.chat_api import _chat_with_openai_rq
# from scripts.server.graph_utils import PyGraph,graph
from tqdm import tqdm


prompt_graph = configs.prompts.prompt_graph

GENERATE_DOMAIN_PROMPT = """
你是一个智能助手，帮助人类分析文本文件中的信息。
给定一个样本文本，请通过分配一个描述性的领域来帮助用户，该领域概括了文本的内容。
示例领域包括：“社会科学”，“算法分析”，“医学科学”等。

文本：{input_text}
"""

GENERATE_PERSONA_PROMPT = """
你是一个人智能助手，帮助人类分析文本文档中的信息。
给定特定类型的任务和示例文本，通过生成一段3到4句话的专家描述来帮助用户解决问题。
使用以下格式：
你是一位专家{{角色}}。你擅长{{相关技能}}。你擅长帮助人们解决{{特定任务}}。

任务：确定实体的关系和结构，特别是在{domain}领域内
角色描述："""

GENERATE_ENTITY_PROMPT = """
这段文字是关于如何从文本中提取实体类型，并以JSON格式返回这些类型。以下是翻译：
目标是研究实体类型及其特征之间的联系和关系，以便理解文本中的所有可用信息。
用户的任务是{task}。
作为分析的一部分，您需要识别以下文本中存在的实体类型。
实体类型必须与用户任务相关。
避免使用“其他”或“未知”等通用实体类型。
这非常重要：不要生成冗余或重叠的实体类型。例如，如果文本中包含“公司”和“组织”实体类型，您应该只返回其中之一。
不用担心数量，总是选择质量而不是数量。并确保您的答案中的每一项都与实体提取的上下文相关。
以JSON格式返回实体类型，使用“entities”作为键，实体类型作为字符串数组。
=====================================================================
示例部分：以下部分包括示例输出。这些示例**必须从您的答案中排除**。
示例1
任务：确定指定社区内的联系和组织层级。
文本：Example_Org_A 是瑞典的一家公司。Example_Org_A 的董事是 Example_Individual_B。
JSON响应：
{"entity_types": ["organization", "person"]}
示例1 结束
示例2
任务：识别不同哲学流派之间共享的关键概念、原则和论点，并追溯它们相互之间的历史或思想影响。
文本：以笛卡尔等思想家为代表的理性主义认为理性是知识的主要来源。该学派的关键概念包括强调演绎推理方法。
JSON响应：
{"entity_types": ["concept", "person", "school of thought"]}
示例2 结束
示例3 任务：识别间接塑造问题的一系列基本力量、因素和趋势。 文本：像松下这样的行业领导者正在争夺电池生产领域的至高无上地位。他们正在大力投资研究与开发，并探索新技术以获得竞争优势。 JSON响应： {"entity_types": ["organization", "technology", "sectors", "investment strategies"]} 示例3 结束
=====================================================================
真实数据：以下部分是真实数据。您应该仅使用此真实数据准备您的答案。仅生成实体类型。
任务：{task}
文本：{input_text}
JSON响应：
{"entity_types": [<entity_types>]}
"""




chunks = get_chunks('./UserData/test.txt',chunk_size = 400,chunk_overlap=30,split_str='\n')



Entities = []
for chunk in chunks:
    input_text = chunk.page_content
    # 领域判定
    query = GENERATE_DOMAIN_PROMPT.format(input_text = input_text)
    response = _chat_with_openai_rq(query,[],0.1)
    domain = response.msg

    # persona判定
    query = GENERATE_DOMAIN_PROMPT.format(domain=domain)
    response = _chat_with_openai_rq(query, [], 0.1)
    persona = response.msg

    # entities判定
    query = GENERATE_ENTITY_PROMPT.format(domain=domain,input_text=input_text)
    response = _chat_with_openai_rq(query,[{'role':'system','content':query}],0.1)
    entities = response.msg

    # 获得Entities
    Entities.append(entities)

# Entities去重
Entities = list(set(Entities))




res = []
for chunk in chunks:
    input_text = chunk.page_content
    # 领域判定
    query = GENERATE_DOMAIN_PROMPT.format(input_text = input_text)
    response = _chat_with_openai_rq(query,[],0.1)


    query = prompt_graph.render(entity_types = Entities, input_text = input_text)
    response = _chat_with_openai_rq(query,[],0.1)
    answer = response.msg
    res.append(answer)


entities, relationships = [], []
for r in res:
    datas = r.split('{record_delimiter}\n')
    for data in datas:
        val = data.split('{tuple_delimiter}')
        if val[0] == '("relationship"':
            relationship,src_entity,tar_entity,relation,strength = val
            relationship, src_entity, tar_entity, relation, strength = relationship.strip('"'), src_entity.strip('"'), tar_entity.strip('"'), relation.strip('"'), strength.strip('"')
            relationships.append([src_entity, tar_entity, relation])  # (源实体，目标实体，实体描述)
        if val[0] == '("entity"':
            entity,entity_name,entity_type,entity_description = val
            entity,entity_name,entity_type,entity_description = entity.strip('"'),entity_name.strip('"'),entity_type.strip('"'),entity_description.strip('"')
            entities.append([entity_name,entity_type,entity_description])     # (实体名称，实体类型，实体描述)


#%% 过滤relationships中未在entities中记录的实体
entities = set([eneity[0] for eneity in entities])
relationship_new = []
for relation in relationships:
    if relation[0] in entities and relation[1] in entities:
        relationship_new.append(relation)

import pickle
with open("relationships.pkl", "wb") as file:
    pickle.dump(relationships, file)

with open("relationships_new.pkl", "wb") as file:
    pickle.dump(relationship_new, file)

with open("entities.pkl", "wb") as file:
    pickle.dump(entities, file)






# 到 neo4j 上
# graph = PyGraph(graph)
#
# for entity_name,entity_type,entity_description in tqdm(entities):
#     graph.create_node(entity_name,entity_type,entity_description)
# for src_entity, tar_entity, relation in tqdm(relationships):
#     graph.create_relation(src_entity,relation,tar_entity)
#
# graph.del_all_graph()



